/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package shaman.storage.impl.filter.published;

import com.thoughtworks.xstream.annotations.XStreamAlias;
import com.thoughtworks.xstream.annotations.XStreamOmitField;
import java.awt.Frame;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInput;
import java.io.DataInputStream;
import java.io.DataOutput;
import java.io.DataOutputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.AlgorithmParameters;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.InvalidParameterException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.Provider;
import java.security.SecureRandom;
import java.security.Security;
import java.security.spec.InvalidKeySpecException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.TreeSet;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKeyFactory;
import shaman.storage.Factory;
import shaman.storage.FilterList;
import shaman.storage.api.Filter;
import shaman.storage.impl.filter.crypto.PasswordKeySpec;
import shaman.storage.impl.util.Bimap;
import shaman.storage.impl.util.StreamUtil;

/**
 *
 * @author dmitriy
 */
@XStreamAlias("crypt")
public class CryptFilter implements Filter {

    public static final Map<String, String> aliases=new Bimap<String>();
    public static final String FILTER_TYPE = "Шифратор";

    public static final String KEY_ALGORITHM_PROPERTY="Алгоритм шифрования ключей";
    public static final String KEY_LENGTH_PROPERTY="Размер ключа шифрования ключей";
    public static final String STREAM_ALGORITHM_PROPERTY="Алгоритм шифрования данных";
    public static final String STREAM_LENGTH_PROPERTY="Размер ключа шифрования данных";
    public static final String ITERATIONS_COUNT_PROPERTY="Число итераций";
    public static final String PASSWORD_PROPERTY="Пароль";

    public static final String PADDING="PKCS5Padding";
    public static final String KEY_CIPHER_MODE="ECB";
    public static final String STREAM_CIPHER_MODE="CBC";
    public static final byte VERSION=1;
    public static final byte SALT_LENGTH=8;
    public static final String ALGORITHM_ENCODING="UTF-8";

    private static class CryptFilterFactory implements Factory<CryptFilter> {

        public String getTypeName() {
            return FILTER_TYPE;
        }

        public CryptFilter newObject() {
            return new CryptFilter();
        }

        public Class<CryptFilter> getGeneratedClass() {
            return CryptFilter.class;
        }

    }
    
    @XStreamAlias("keyAlgorithm")
    private String keyAlgorithm="AES";

    @XStreamAlias("keySize")
    private int keySize=256;

    @XStreamAlias("streamAlgorithm")
    private String streamAlgorithm="AES";
    
    @XStreamAlias("streamKeySize")
    private int streamKeySize=256;
    
    @XStreamAlias("iterationCount")
    private int iterationCount=1000;

    @XStreamOmitField
    private String password;

    public CryptFilter() {
    }

    public FilterType getFilterType() {
        return FilterType.CRYPTOR;
    }

    public String getName() {
        return FILTER_TYPE;
    }

    public OutputStream getOutputStream(OutputStream cipheredStream) throws IOException {
        if (password==null)
            throw new IOException("Не указан пароль");
        Cipher keyCipher;
        Cipher cipher;
        byte[] randomKey;
        byte[] salt;
        byte[] params;
        Key streamKey;
        try {
            KeyGenerator keyGen=KeyGenerator.getInstance(streamAlgorithm);
            keyGen.init(streamKeySize);
            streamKey=keyGen.generateKey();
            keyCipher=Cipher.getInstance(keyAlgorithm+"/"+KEY_CIPHER_MODE+"/"+PADDING);
            salt=new byte[SALT_LENGTH];
            new SecureRandom().nextBytes(salt);
            Key secretKey=SecretKeyFactory.getInstance(keyAlgorithm).generateSecret(new PasswordKeySpec(password, salt, keySize/8, iterationCount, keyAlgorithm));
            keyCipher.init(Cipher.WRAP_MODE, secretKey);
            cipher=Cipher.getInstance(streamAlgorithm+"/"+STREAM_CIPHER_MODE+"/"+PADDING);
            cipher.init(Cipher.ENCRYPT_MODE, streamKey);
            params=cipher.getParameters().getEncoded();
            randomKey=keyCipher.wrap(streamKey);
        } catch (IllegalBlockSizeException ex) {
            throw new IOException("Неверный размер блока", ex);
        } catch (NoSuchAlgorithmException ex) {
            throw new IOException("Алгоритм шифрования недоступен", ex);
        } catch (NoSuchPaddingException ex) {
            throw new IOException("Алгоритм выравнивания "+PADDING+" недоступен", ex);
        } catch (InvalidKeyException ex) {
            throw new IOException("Неверные параметры ключа", ex);
        } catch (InvalidKeySpecException ex) {
            throw new IOException("Не удалось сгенерировать ключ", ex);
        }
        cipheredStream.write(VERSION);
        StreamUtil.writeString(cipheredStream, keyAlgorithm);
        StreamUtil.writeUtf8Char(cipheredStream, keySize);
        StreamUtil.writeBlob(cipheredStream, salt);
        StreamUtil.writeUtf8Char(cipheredStream, iterationCount);
        StreamUtil.writeString(cipheredStream, streamAlgorithm);
        StreamUtil.writeBlob(cipheredStream, randomKey);
        StreamUtil.writeBlob(cipheredStream, params);
        return new CipherOutputStream(cipheredStream, cipher);
    }

    public InputStream getInputStream(InputStream cipheredStream) throws IOException {
        if (password==null)
            throw new IOException("Не указан пароль");
        Cipher keyCipher;
        Cipher cipher;
        Key streamKey;
        String curKeyAlgorithm, curStreamAlgorithm;
        byte[] randomKey;
        byte[] salt;
        byte[] params;
        int curIterationCount, curKeySize;
        int version=cipheredStream.read();
        switch(version) {
            case 1:
                curKeyAlgorithm=StreamUtil.readString(cipheredStream);
                curKeySize=StreamUtil.readUtf8Char(cipheredStream);
                salt=StreamUtil.readBlob(cipheredStream);
                curIterationCount=StreamUtil.readUtf8Char(cipheredStream);
                curStreamAlgorithm=StreamUtil.readString(cipheredStream);
                randomKey=StreamUtil.readBlob(cipheredStream);
                params=StreamUtil.readBlob(cipheredStream);
            break;
            case -1:
                throw new IOException("Неожиданный конец потока");
            default:
                throw new IOException("Неизвестная версия формата файла");
        }
        try {
            keyCipher=Cipher.getInstance(curKeyAlgorithm+"/"+KEY_CIPHER_MODE+"/"+PADDING);
            Key secretKey=new PasswordKeySpec(password, salt, curKeySize/8, curIterationCount, curKeyAlgorithm);
            keyCipher.init(Cipher.UNWRAP_MODE, secretKey);
            cipher=Cipher.getInstance(curStreamAlgorithm+"/"+STREAM_CIPHER_MODE+"/"+PADDING);
            streamKey=keyCipher.unwrap(randomKey, curStreamAlgorithm, Cipher.SECRET_KEY);
            AlgorithmParameters algParams=AlgorithmParameters.getInstance(streamAlgorithm);
            algParams.init(params);
            cipher.init(Cipher.DECRYPT_MODE, streamKey, algParams);
        } catch (NoSuchAlgorithmException ex) {
            throw new IOException("Алгоритм шифрования недоступен", ex);
        } catch (NoSuchPaddingException ex) {
            throw new IOException("Алгоритм выравнивания "+PADDING+" недоступен", ex);
        } catch (InvalidKeyException ex) {
            throw new IOException("Неверные параметры ключа", ex);
        } catch (InvalidAlgorithmParameterException ex) {
            throw new IOException("Неверные параметры алгоритма", ex);
        }
        return new CipherInputStream(cipheredStream, cipher);
    }

    public boolean hasConfigForm() {
        return false;
    }

    public boolean showConfigForm(Frame parent) {
        throw new UnsupportedOperationException("Not supported yet.");
    }

    public List<String> getPropertyNames() {
        return Arrays.asList(new String[]{
            KEY_ALGORITHM_PROPERTY,
            KEY_LENGTH_PROPERTY,
            STREAM_ALGORITHM_PROPERTY,
            STREAM_LENGTH_PROPERTY,
            ITERATIONS_COUNT_PROPERTY,
            PASSWORD_PROPERTY
        });
    }

    public String getPropertyValue(String name) throws InvalidParameterException {
        if (name.equals(KEY_ALGORITHM_PROPERTY))
            return keyAlgorithm;
        else if (name.equals(KEY_LENGTH_PROPERTY))
            return String.valueOf(keySize);
        else if (name.equals(STREAM_ALGORITHM_PROPERTY))
            return streamAlgorithm;
        else if (name.equals(STREAM_LENGTH_PROPERTY))
            return String.valueOf(streamKeySize);
        else if (name.equals(ITERATIONS_COUNT_PROPERTY))
            return String.valueOf(iterationCount);
        else if (name.equals(PASSWORD_PROPERTY))
        {
            if (password==null || password.length()==0)
                return "";
            else
                return "********";
        }
        else
            throw new InvalidParameterException("Неизвестное имя свойства: "+name);
    }

    public static Set<String> getCipherAlgorithms() {
        Set<String> algs=new TreeSet<String>();
        for (Provider p: Security.getProviders())
        {
            for (Provider.Service s: p.getServices())
            {
                if (s.getType().equals("Cipher"))
                {
                    String alg=s.getAlgorithm();
                    try {
                        KeyGenerator.getInstance(alg, p);
                        String altname=aliases.get(alg.toUpperCase());
                        if (altname==null)
                            algs.add(alg);
                        else
                            algs.add(alg+" ("+altname+")");
                    } catch (NoSuchAlgorithmException ex) {
                    }
                }
            }
        }
        return algs;
    }

    public static boolean checkKeySize(String algorithm, int size)
    {
        try {
            KeyGenerator.getInstance(algorithm).init(size);
            return true;
        } catch (NoSuchAlgorithmException ex) {
        } catch (InvalidParameterException ex) {
        }
        return false;
    }

    public static Set<Integer> getKeySizes(String algorithm) {
        Set<Integer> sizes=new TreeSet<Integer>();
        int i;
        KeyGenerator kgen;
        Random rand=new Random();
        Matcher numbers=null;
        Pattern p=Pattern.compile("[0-9][0-9]*");
        try {
            kgen=KeyGenerator.getInstance(algorithm);
        } catch (NoSuchAlgorithmException ex) {
            return null;
        }
        for (i=0; i<10; i++)
        {
            int size=rand.nextInt(50000);
            try {
                kgen.init(size);
            } catch (InvalidParameterException ex) {
                String msg=ex.getMessage();
                Matcher m=p.matcher(msg);
                if (m.find())
                {
                    int j;
                    for (m.reset(); m.find(); )
                    {
                        String part=m.group();
                        int n;
                        try {
                            n=Integer.parseInt(part);
                            if (n!=0 && n!=size)
                            {
                                numbers=m;
                                break;
                            }
                        } catch (NumberFormatException exx) {
                        }
                    }
                }
            }
            if (numbers!=null)
                break;
        }
        if (numbers==null)
            return null;
        for (numbers.reset(); numbers.find(); )
        {
            String num=numbers.group();
            try {
                int size=Integer.parseInt(num);
                if (checkKeySize(algorithm, size))
                    sizes.add(size);
            } catch (NumberFormatException ex) {
            }
        }
        if (!sizes.isEmpty())
            return sizes;
        else
            return null;
    }

    public static List<String> getSizesList(String name) {
        Set<Integer> sizes=getKeySizes(name);
        if (sizes==null)
            return null;
        List<String> sizeStrings=new ArrayList<String>();
        for (Integer size: sizes)
            sizeStrings.add(size.toString());
        return sizeStrings;
    }

    public List<String> getPropertyValues(String name) throws InvalidParameterException {
        if (name.equals(KEY_ALGORITHM_PROPERTY) || name.equals(STREAM_ALGORITHM_PROPERTY))
        {
            return new ArrayList<String>(getCipherAlgorithms());
        }
        else if (name.equals(KEY_LENGTH_PROPERTY))
        {
            return getSizesList(keyAlgorithm);
        }
        else if (name.equals(STREAM_LENGTH_PROPERTY))
        {
            return getSizesList(streamAlgorithm);
        }
        else
            return null;
    }

    private boolean checkAlgorithm(String name, String mode, String padding) {
        try {
            KeyGenerator.getInstance(name);
            Cipher.getInstance(name+"/"+mode+"/"+padding);
            return true;
        } catch (NoSuchAlgorithmException ex) {
        } catch (NoSuchPaddingException ex) {
        }
        try {
            if (Cipher.getInstance(name).getBlockSize()==0)
                return true;
        } catch (NoSuchAlgorithmException ex) {
        } catch (NoSuchPaddingException ex) {
        }
        return false;
    }

    private String getAlhorithmName(String name, String mode, String padding) throws InvalidParameterException {
        if (checkAlgorithm(name, mode, padding))
            return name;
        String newName=aliases.get(name.toUpperCase());
        if (newName!=null && checkAlgorithm(newName, mode, padding))
            return newName;
        Matcher m;
        m=Pattern.compile("^([^(]*) [(]([^)]*)[)]$").matcher(name);
        if (!m.matches())
            throw new InvalidParameterException("Неизвестный или неподходящий алгоритм шифрования");
        String name1=m.group(1);
        String name2=m.group(2);
        String regName2=aliases.get(name1.toUpperCase());
        if (regName2==null)
            throw new InvalidParameterException("Неизвестный или неподходящий алгоритм шифрования");
        if (name2.toUpperCase().equalsIgnoreCase(name2))
        {
            if (checkAlgorithm(name1, mode, padding))
                return name1;
            else if (checkAlgorithm(name2, mode, padding))
                return name2;
        }
        throw new InvalidParameterException("Неизвестный или неподходящий алгоритм шифрования");
    }

    public void setPropertyValue(String name, String value) throws InvalidParameterException {
        if (name.equals(KEY_ALGORITHM_PROPERTY))
            keyAlgorithm=getAlhorithmName(value, KEY_CIPHER_MODE, PADDING);
        else if (name.equals(KEY_LENGTH_PROPERTY))
        {
            int val;
            try {
                val=Integer.parseInt(value);
            } catch (NumberFormatException ex) {
                throw new InvalidParameterException("Размер ключа должен быть числом");
            }
            if (val<=0)
                throw new InvalidParameterException("Размер ключа должен быть положительным");
            keySize=val;
        }
        else if (name.equals(STREAM_ALGORITHM_PROPERTY))
            streamAlgorithm=getAlhorithmName(value, STREAM_CIPHER_MODE, PADDING);
        else if (name.equals(STREAM_LENGTH_PROPERTY))
        {
            int val;
            try {
                val=Integer.parseInt(value);
            } catch (NumberFormatException ex) {
                throw new InvalidParameterException("Размер ключа должен быть числом");
            }
            if (val<=0)
                throw new InvalidParameterException("Размер ключа должен быть положительным");
            streamKeySize=val;
        }
        else if (name.equals(ITERATIONS_COUNT_PROPERTY))
        {
            int val;
            try {
                val=Integer.parseInt(value);
            } catch (NumberFormatException ex) {
                throw new InvalidParameterException("Число итераций должно быть числом");
            }
            if (val<=0)
                throw new InvalidParameterException("Число итераций должно быть положительным");
            iterationCount=val;
        }
        else if (name.equals(PASSWORD_PROPERTY))
        {
            password=value;
        }
        else
            throw new InvalidParameterException("Неизвестное имя свойства: "+name);
    }

    static {
        aliases.put("DESEDE", "3DES");
        aliases.put("ARCFOUR", "RC4");
        FilterList.addFactory(new CryptFilterFactory());
    }

    public static void main(String[] args)
    {
        try {
            try {
                Cipher.getInstance("AES/KBT/none");
            } catch (Exception ex) {
                ex.printStackTrace();
            }
            Filter cryptFilter=FilterList.getFactoryMap().get(FILTER_TYPE).newObject();
            for (String s: cryptFilter.getPropertyValues(STREAM_ALGORITHM_PROPERTY))
            {
                System.out.println(s);
            }
            new DeflateFilter();
            cryptFilter.setPropertyValue(PASSWORD_PROPERTY, "AVerySecretLongAndStrongPassword");
            cryptFilter.setPropertyValue(STREAM_ALGORITHM_PROPERTY, "ARCFOUR");
            for (String s: cryptFilter.getPropertyValues(STREAM_LENGTH_PROPERTY))
            {
                System.out.println(s);
            }
            cryptFilter.setPropertyValue(STREAM_LENGTH_PROPERTY, "56");
            ByteArrayOutputStream encryptedStream=new ByteArrayOutputStream();
            OutputStream encryptStream=cryptFilter.getOutputStream(encryptedStream);
            DataOutput encryptData=new DataOutputStream(encryptStream);
            encryptData.writeUTF("Hello World! Hello World! Hello World!");
            encryptStream.close();
            encryptedStream.close();
            byte[] encryptedData=encryptedStream.toByteArray();
            StringBuilder dataHex=new StringBuilder("Encrypted:");
            int i;
            for (i=0; i<encryptedData.length; i++)
            {
                dataHex.append(" ");
                String hex=Integer.toHexString((encryptedData[i]+256)%256);
                if (hex.length()<2)
                    dataHex.append("0");
                dataHex.append(hex);
            }
            OutputStream fileOut=new FileOutputStream("encrypted.bin");
            fileOut.write(encryptedData);
            fileOut.close();
            System.out.println(dataHex.toString());
            InputStream decryptStream=new ByteArrayInputStream(encryptedData);
            InputStream decryptedStream=cryptFilter.getInputStream(decryptStream);
            DataInput decryptData=new DataInputStream(decryptedStream);
            String stringRead=decryptData.readUTF();
            decryptedStream.close();
            decryptStream.close();
            System.out.println(stringRead);
        } catch (IOException ex) {
            ex.printStackTrace();
        }
    }
}
